import collections
import datetime
import pandas
import os
import numpy
import psutil
import resource
import tracemalloc
from inspect import getframeinfo, stack

import coin.proto.archive_flag_pb2 as afpb

from coin.base import flags
from coin.exchange.kr_rest.product.product_impl import generate_product_from_str2
from coin.exchange.util.feed_checker import get_default_feed_checker_config

from coin.strategy.mm.dumper_base import run_from_archive
from coin.strategy.mm.subscription import get_subscription_key
from coin.strategy.mm.subscription import FeedSubscriptionRequest
from coin.support.proto_log.logic.util import (run_from_strat_log_archive, run_from_og_log_archive)

from coin.proto.coin_order_gateway_pb2 import OrderGatewayLog
import coin.feed.fastfeed.native_wrapper as nwrp


def print_mem(obj=None):
  if obj is not None:
    import pympler.asizeof
    print(pympler.asizeof.asizeof(obj))
  else:
    caller = getframeinfo(stack()[1][0])
    print("----------")
    print("%s:%d" % (caller.filename, caller.lineno))
    print("ru_maxrss: {:.2f}MiB".format(
          resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024))
    print("virtual_memory.used: {:.2f}MiB".format(
          psutil.virtual_memory().used / 1024 ** 2))
    print("----------")


def profile_mem():
  import guppy
  print(guppy.hpy().heap())


def trace_memory_begin():
  tracemalloc.start()


def trace_memory_print():
  snapshot = tracemalloc.take_snapshot()
  top_stats = snapshot.statistics('lineno')
  print("[ Top 10 ]")
  for stat in top_stats[:10]:
    print(stat)


def define_base_flags():
  flags.DEFINE_string('trading_date', None, 'Trading date in form of %Y%m%d.')


def define_feed_archive_flags():
  flags.DEFINE_string('exchange', None, 'Exchange name.')

  flags.DEFINE_string('market_type', None, 'Spot, Futures')

  flags.DEFINE_string('feed_machine', None, 'feed machine name.')

  flags.DEFINE_string('feed_root', '/remote/iosg/coin/data/flow', 'feed machine name.')

  flags.DEFINE_string('time_range', '0-24', '')

  flags.DEFINE_string('worker_ids', '0', '')

  flags.DEFINE_string('api_override', None, '')

  flags.DEFINE_bool('use_feed_cache', True, '')

  flags.DEFINE_bool('use_swap_over_futures', False, '')

  flags.DEFINE_string('owner', None, '')

  flags.DEFINE_string('feed_cache_dir', '/remote/iosg/coin-2/feed_cache', '')

  flags.DEFINE_boolean('include_funding_rate', False, '')

  flags.DEFINE_boolean('bbo_from_cache', False, '')

  flags.DEFINE_boolean('use_pickle', False, '')

  flags.DEFINE_string(
      'bbo_cache_root', '/remote/iosg/data-2/buckets/feed.derived.interval_h5/coin/main/PT1M/', '')


def define_order_archive_flags():
  flags.DEFINE_string('orderlog_machine', None, 'Instance machine name.')

  flags.DEFINE_string('orderlog_root', '/remote/iosg/strat-1/buckets/log.raw.coin/live/og_proto_log', '')

  flags.DEFINE_string('strategy_name', None, 'Strategy Name')

  flags.DEFINE_string('stratlog_cache_root', '/remote/iosg/data-2/buckets/order.derived.interval_h5/coin/PT1M', '')


def get_base_flags():
  assert flags.FLAGS.trading_date is not None, flags.FLAGS.trading_date

  return afpb.BaseFlagProto(
      trading_date=flags.FLAGS.trading_date,
      exchange=flags.FLAGS.exchange,
      market_type=flags.FLAGS.market_type,
      time_range=flags.FLAGS.time_range,
  )


def split_time_range(time_range, split_hours):
  start, end = time_range.split("-")
  start = float(start)
  end = float(end)
  timenodes = numpy.append(numpy.arange(start, end, split_hours), end)
  return [
      f"{'%05.2f' % st}-{'%05.2f' % fi}"
      for st, fi in zip(timenodes[:-1], timenodes[1:])]


def get_api_override_from_override_string(api_override):
  if api_override is None or api_override == "":
    return {}
  return {value.split("=")[0]: value.split("=")[1] for value in api_override.split(",")}


def is_coin2_machine(feed_machine):
  ret = ((feed_machine.find("feed-05") >= 0) or feed_machine == "feed-best.ap-northeast-1.aws"
         or feed_machine == "feed-best.cn-hongkong.aliyun"
         or feed_machine in ["feed-02.eu-west-1.aws", "feed-02.us-east-1.aws"])
  return ret


def get_default_api_override(coin2):
  if flags.FLAGS.use_swap_over_futures:
    addstr = (
        "Futures:Okex=v5-swap.realtime_trade_move2bp,"
        "Futures:Huobi=v1-swap.bbo_trade_move2bp,")
  else:
    addstr = (
        "Futures:Okex=v5.realtime_trade_move2bp,"
        "Futures:Huobi=v1.bbo_trade_move2bp,")

  return get_api_override_from_override_string(addstr +
                                               "Futures:Binance=v1.l1_realtime_trade_move2bp,"
                                               "Spot:Binance=v1.l1_realtime_trade_move2bp,"
                                               "Spot:Gdax=v1.realtime_move2bp,"
                                               "Spot:Okex=v5.realtime_move2bp,"
                                               "Spot:Huobi=v1.bbo_trade_move2bp,"
                                               "Futures:Ftx=v1.bbo_trade_move2bp,"
                                               "Futures:Bitmex=v1.realtime_light_bbo,"
                                               "Spot:Quoinex=v2.realtime,"
                                               "Futures:Ftx=v1.bbo_trade_move2bp,"
                                               "Spot:Ftx=v1.bbo_trade_move2bp,"
                                               "Spot:Bithumb=v2.realtime,"
                                               "Spot:Coinone=v2.realtime,"
                                               "Spot:Upbit=v1.realtime,"
                                               "Spot:Kraken=v1.snapshot,"
                                               "Spot:Bitstamp=v2.snapshot,"
                                               "Spot:Gateio=v4.bbo,"
                                               "Spot:Mexc=v3.realtime_light_bbo,"
                                               "Spot:Kucoin=v1.realtime,"
                                               "Futures:Bitflyer=v1.realtime,"
                                               "Futures:Bybit=v2.realtime,"
                                               "Spot:Mercado=v4.realtime,"
                                               "Spot:Bitbank=v1.realtime")


def get_api_override_from_flags(feed_machine):
  if flags.FLAGS.api_override is None:
    return None
  else:
    return get_api_override_from_override_string(flags.FLAGS.api_override)


def get_feed_subscription_flags():
  assert flags.FLAGS.feed_machine is not None, flags.FLAGS.feed_machine
  return afpb.FeedSubscriptionFlagProto(feed_machine=flags.FLAGS.feed_machine,
                                        feed_root=flags.FLAGS.feed_root,
                                        worker_ids=flags.FLAGS.worker_ids,
                                        api_override=get_api_override_from_flags(
                                            flags.FLAGS.feed_machine),
                                        use_feed_cache=flags.FLAGS.use_feed_cache,
                                        feed_cache_dir=os.path.join(flags.FLAGS.feed_cache_dir))


def get_order_subscription_flags():
  return afpb.OrderSubscriptionFlagProto(orderlog_machine=flags.FLAGS.orderlog_machine,
                                         orderlog_root=os.path.expanduser(
                                             flags.FLAGS.orderlog_root),
                                         owner=flags.FLAGS.owner,
                                         strategy_name=flags.FLAGS.strategy_name)


def get_time_range_from_trading_date_str(trading_date, time_range):
  if trading_date.find("-") >= 0:
    from_date, end_date = trading_date.split("-")
  else:
    from_date = trading_date
    end_date = trading_date
  add_hour, hours = time_range.split("-")
  add_hour = float(add_hour)
  hours = float(hours) - add_hour
  from_date = datetime.datetime.strptime(from_date, '%Y%m%d')
  end_date = datetime.datetime.strptime(end_date, '%Y%m%d')
  start_time = from_date + datetime.timedelta(hours=add_hour)
  end_time = end_date + datetime.timedelta(hours=hours) + datetime.timedelta(hours=add_hour)
  return start_time, end_time


def get_time_range(baseflags):
  return get_time_range_from_trading_date_str(baseflags.trading_date, baseflags.time_range)


def symbols_to_products(symbols, baseflags):
  add_hour, hours = baseflags.time_range.split("-")
  start_time, _ = get_time_range(baseflags)

  products = [
      generate_product_from_str2(market_type=baseflags.market_type,
                                 exchange=baseflags.exchange,
                                 api_version=None,
                                 norm_product=symbol,
                                 current_datetime=start_time) for symbol in symbols
  ]
  return products


# Okex:Futures:BTC-USD.QUARTER


def subsymbols_to_products(subsymbols, baseflags):
  add_hour, hours = baseflags.time_range.split("-")
  start_time, _ = get_time_range(baseflags)

  products = []
  for subsymbol in subsymbols:
    splits = subsymbol.split(":")
    assert len(splits) == 3, splits
    exchange, market_type, symbol = splits[0], splits[1], splits[2]

    products.append(
        generate_product_from_str2(market_type=market_type,
                                   exchange=exchange,
                                   api_version=None,
                                   norm_product=symbol,
                                   current_datetime=start_time))
  return products


def get_feed_machine_override_map(feed_sub_request, machine):
  override = None
  if machine in ["feed-coin1.ap-northeast-1.aws", "feed-best.ap-northeast-1.aws"]:
    override = {}
    for sub_req, products in feed_sub_request.gen_iter():
      if sub_req.to_str().count(".") >= 3:
        override[sub_req] = "feed-05.ap-northeast-1.aws"
      else:
        override[sub_req] = "feed-04.ap-northeast-1.aws"
  elif machine in ["feed-coin1.cn-hongkong.aliyun", "feed-best.cn-hongkong.aliyun"]:
    override = {}
    for sub_req, products in feed_sub_request.gen_iter():
      if sub_req.to_str().count(".") >= 3:
        override[sub_req] = "feed-05.cn-hongkong.aliyun"
      else:
        override[sub_req] = "feed-04.cn-hongkong.aliyun"
  elif machine.startswith("feed-best.ap-northeast-2.aws"):
    override = {}
    for sub_req, products in feed_sub_request.gen_iter():
      if sub_req.to_str().count(".") >= 3:
        override[sub_req] = "feed-02.ap-northeast-2.aws"
      else:
        override[sub_req] = "feed-01.ap-northeast-2.aws"
  elif machine.startswith("feed-best"):
    override = {}
    for sub_req, products in feed_sub_request.gen_iter():
      override[sub_req] = machine.replace("feed-best", "feed-01")
  return override


def get_product_api_version(product):
  if product.exchange in ["Bithumb", "Bitstamp", "Coinone", "Quoinex", "Maicoin"]:
    return "v2"
  elif product.exchange in ["Mexc", "Bitmart"]:
    return "v3"
  elif product.exchange in ["Okex"]:
    return "v5"
  elif product.exchange in ["Gateio", "Mercado"]:
    return "v4"
  else:
    return "v1"


def get_coin2_recipe(mea_coin1):
  if ".Ftx" in mea_coin1:
    return "bbo_trade_move2bp"
  if "Spot.Okex.v5" in mea_coin1:
    return "realtime_trade_move2bp"
  elif "Futures.Okex.v5" in mea_coin1:
    return "realtime_trade_move2bp"
  elif "Futures.Huobi" in mea_coin1:
    return "bbo_move2bp"
  elif "Futures.Bitmex" in mea_coin1:
    return "realtime_light_bbo"
  elif "Spot.Binance" in mea_coin1:
    return "l1_realtime_trade_move2bp"
  elif "Futures.Binance" in mea_coin1:
    return "l1_realtime_trade_move2bp"
  elif "Spot.Huobi" in mea_coin1:
    return "bbo"
  elif "Spot.Gdax" in mea_coin1:
    return "realtime_trade_move2bp"
  elif "Gateio" in mea_coin1:
    return "bbo"
  elif "Bitmart" in mea_coin1:
    return "snapshot"
  elif "Spot.Mexc" in mea_coin1:
    return "realtime_light_bbo"
  else:
    return "realtime"


def get_recipe_light(mea):
  mea_coin1 = ".".join(mea)
  if "Binance" in mea_coin1:
    return "l1_realtime_move2bp"
  elif "Futures.Huobi" in mea_coin1:
    return "bbo_move2bp"
  elif (
      "Futures.Okex" in mea_coin1 or
      "Gdax" in mea_coin1):
    return "realtime_move2bp"
  else:
    return get_coin2_recipe(mea_coin1)


Mear = collections.namedtuple("Mear", ["market", "exchange", "api_version", "recipe"])

def get_mear_light(product):
  mea = [product.market_name, product.exchange, get_product_api_version(product)]
  return Mear(*mea, get_recipe_light(mea))


def run_from_archive_product(
    *,
    trading_date,
    time_range,
    feed_machine,
    product,
    api_version,
    recipe,
    on_book_reset):
  start_time, end_time = get_time_range_from_trading_date_str(trading_date, time_range)
  feed_sub_request = FeedSubscriptionRequest()
  market = product.market_name
  exch = product.exchange

  api_version_override_map = {(market, exch): f"{api_version}.{recipe}"}
  feed_sub_request.add_products(
      [product],
      get_subscription_key(product, api_version_override_map=api_version_override_map))
  start_trading_date = start_time.date()
  end_trading_date = end_time.date()
  begin_of_date = datetime.datetime.combine(
      start_trading_date, datetime.datetime.min.time())
  end_of_date = datetime.datetime.combine(
      end_trading_date + datetime.timedelta(days=1),
      datetime.datetime.min.time())
  nwrp.run_from_feed_cache(
      feed_sub_request,
      on_book_reset,
      max(start_time - datetime.timedelta(hours=0.5), begin_of_date),
      min(end_time + datetime.timedelta(hours=0.5), end_of_date),
      machine=feed_machine,
      worker_id='0',
      feed_cache_dir='/remote/iosg/coin-2/feed_cache',
      sub_req_to_machine_override_map=get_feed_machine_override_map(
          feed_sub_request, feed_machine))


def run_from_archive_norm(
    *,
    trading_date,
    time_range,
    norm_symbol,
    feed_machine,
    on_book_reset):
  market, exch, api_version, recipe, pname = norm_symbol.split(":")
  start_time, end_time = get_time_range_from_trading_date_str(trading_date, time_range)

  product = generate_product_from_str2(
      market_type=market,
      exchange=exch,
      api_version=None,
      norm_product=pname,
      current_datetime=start_time)
  yield product
  return run_from_archive_product(
      trading_date=trading_date,
      time_range=time_range,
      product=product,
      feed_machine=feed_machine,
      on_book_reset=on_book_reset)


def run_from_feed_cache_with_feed_sub(feed_sub_request,
                                      on_book_reset,
                                      start_time,
                                      end_time,
                                      machine,
                                      worker_ids,
                                      *,
                                      sub_req_hack_map=None,
                                      feed_cache_dir=None):
  feed_cache_dir = feed_cache_dir or '/remote/iosg/coin-2/feed_cache'

  trading_date = start_time.date()
  begin_of_date = datetime.datetime.combine(trading_date, datetime.datetime.min.time())
  end_of_date = datetime.datetime.combine(trading_date + datetime.timedelta(days=1),
                                          datetime.datetime.min.time())

  nwrp.run_from_feed_cache(feed_sub_request,
                           on_book_reset,
                           max(start_time - datetime.timedelta(hours=0.5), begin_of_date),
                           min(end_time + datetime.timedelta(hours=0.5), end_of_date),
                           sub_req_hack_map=sub_req_hack_map,
                           machine=machine,
                           worker_id=worker_ids,
                           feed_cache_dir=feed_cache_dir,
                           sub_req_to_machine_override_map=get_feed_machine_override_map(
                               feed_sub_request, machine))


def products_to_feedsub_request(products, api_version_override_map=None):
  feed_sub_request = FeedSubscriptionRequest()
  for product in products:
    feed_sub_request.add_products(
        [product],
        get_subscription_key(
            product, api_version_override_map=api_version_override_map))
  return feed_sub_request


def products_to_feedsub_request_with_api_override(products, api_override, feed_machine):
  if api_override is None or len(api_override) == 0:
    api_override = get_default_api_override(is_coin2_machine(feed_machine))
  api_version_override_map = {
      tuple(key.split(":")): (None if value == "" else value) for key,
      value in api_override.items()
  }
  return products_to_feedsub_request(products, api_version_override_map)


def run_from_feed_cache_base(products,
                             on_book_reset,
                             start_time,
                             end_time,
                             machine,
                             worker_ids,
                             *,
                             feed_cache_dir=None,
                             api_override=None,
                             write_feed_cache_if_none=False):
  run_from_feed_cache_with_feed_sub(products_to_feedsub_request_with_api_override(
      products, api_override, machine),
                                    on_book_reset,
                                    start_time,
                                    end_time,
                                    machine,
                                    worker_ids,
                                    feed_cache_dir=feed_cache_dir)


def run_from_archive_base_with_feed_sub(baseflags,
                                        feed_sub_flags_proto,
                                        *,
                                        feed_sub_request,
                                        on_book_reset):
  start_time, end_time = get_time_range(baseflags)
  worker_ids = feed_sub_flags_proto.worker_ids.split(",")

  products = feed_sub_request.all_products
  feed_checker_config = get_default_feed_checker_config(products)
  if feed_sub_flags_proto.use_feed_cache:
    run_from_feed_cache_with_feed_sub(feed_sub_request,
                                      on_book_reset,
                                      start_time,
                                      end_time,
                                      feed_sub_flags_proto.feed_machine,
                                      worker_ids,
                                      feed_cache_dir=feed_sub_flags_proto.feed_cache_dir)
  else:
    trading_date = start_time.date()
    begin_of_date = datetime.datetime.combine(trading_date, datetime.datetime.min.time())
    end_of_date = datetime.datetime.combine(trading_date + datetime.timedelta(days=1),
                                            datetime.datetime.min.time())

    run_from_archive(feed_sub_request,
                     on_book_reset,
                     max(start_time - datetime.timedelta(hours=0.5), begin_of_date),
                     min(end_time + datetime.timedelta(hours=0.5), end_of_date),
                     machine=[feed_sub_flags_proto.feed_machine],
                     root_dir=feed_sub_flags_proto.feed_root
                     if feed_sub_flags_proto.HasField('feed_root') else None,
                     worker_id=worker_ids[0] if len(worker_ids) == 1 else worker_ids,
                     feed_checker_config=feed_checker_config,
                     use_feed_arbitration=(len(worker_ids) > 1),
                     sub_req_to_machine_override_map=get_feed_machine_override_map(
                         feed_sub_request, feed_sub_flags_proto.feed_machine))


def run_from_archive_base(baseflags, feed_sub_flags_proto, *, products, on_book_reset):
  start_time, end_time = get_time_range(baseflags)
  worker_ids = feed_sub_flags_proto.worker_ids.split(",")

  feed_checker_config = get_default_feed_checker_config(products)

  if feed_sub_flags_proto.use_feed_cache:
    run_from_feed_cache_base(products,
                             on_book_reset,
                             start_time,
                             end_time,
                             feed_sub_flags_proto.feed_machine,
                             worker_ids[0] if len(worker_ids) == 1 else worker_ids,
                             write_feed_cache_if_none=feed_sub_flags_proto.write_feed_cache_if_none,
                             feed_cache_dir=feed_sub_flags_proto.feed_cache_dir,
                             api_override=feed_sub_flags_proto.api_override)
  else:
    trading_date = start_time.date()
    begin_of_date = datetime.datetime.combine(trading_date, datetime.datetime.min.time())
    end_of_date = datetime.datetime.combine(trading_date + datetime.timedelta(days=1),
                                            datetime.datetime.min.time())
    feed_sub_request = products_to_feedsub_request_with_api_override(
        products, feed_sub_flags_proto.api_override, feed_sub_flags_proto.feed_machine)

    run_from_archive(feed_sub_request,
                     on_book_reset,
                     max(start_time - datetime.timedelta(hours=0.5), begin_of_date),
                     min(end_time + datetime.timedelta(hours=0.5), end_of_date),
                     machine=[feed_sub_flags_proto.feed_machine],
                     root_dir=feed_sub_flags_proto.feed_root
                     if feed_sub_flags_proto.HasField('feed_root') else None,
                     worker_id=worker_ids[0] if len(worker_ids) == 1 else worker_ids,
                     feed_checker_config=feed_checker_config,
                     use_feed_arbitration=(len(worker_ids) > 1),
                     sub_req_to_machine_override_map=get_feed_machine_override_map(
                         feed_sub_request, feed_sub_flags_proto.feed_machine))


def run_from_strat_archive_base(baseflags, order_flags_proto, *, on_log):
  start_time, end_time = get_time_range(baseflags)

  run_from_strat_log_archive(on_log_callback=on_log,
                             start_time=start_time - datetime.timedelta(hours=2.0),
                             end_time=end_time + datetime.timedelta(hours=2.0),
                             machine=order_flags_proto.orderlog_machine,
                             root_dir="/remote/iosg/strat-1/buckets/log.raw.coin/live/strat_proto_log",
                             strategy_name=order_flags_proto.strategy_name)


def run_from_order_archive_base(baseflags, order_flags_proto, *, on_log):
  start_time, end_time = get_time_range(baseflags)

  run_from_og_log_archive(
      on_log,
      baseflags.market_type,
      baseflags.exchange,
      start_time - datetime.timedelta(hours=2.0),
      end_time + datetime.timedelta(hours=2.0),
      machine=order_flags_proto.orderlog_machine,
      root_dir=order_flags_proto.orderlog_root,
      owner=order_flags_proto.owner if len(order_flags_proto.owner) > 0 else None)


def run_from_stateless_order_archive(on_log,
                                     market_type,
                                     exchange,
                                     machine,
                                     root_dir,
                                     owner,
                                     start_time: datetime.datetime,
                                     end_time: datetime.datetime):
  run_from_og_log_archive(on_log,
                          market_type,
                          exchange,
                          start_time - datetime.timedelta(hours=2.0),
                          end_time + datetime.timedelta(hours=2.0),
                          machine=machine,
                          root_dir=root_dir,
                          owner=owner)


class GetAllSymbolsOnLogStrategy(object):
  def __init__(self):
    self.symbols = set()

  def on_log(self, timestamp, log):
    pb = OrderGatewayLog()
    try:
      pb.ParseFromString(log)
      if pb.type == OrderGatewayLog.ORDER_EVENT:
        self.symbols.add(pb.event.symbol)
    except Exception:
      return

  def get_all_symbols(self):
    return list(self.symbols)


def get_all_symbols_from_og_archive(baseflags, order_flags_proto):
  all_symbols_strat = GetAllSymbolsOnLogStrategy()
  run_from_order_archive_base(baseflags, order_flags_proto, on_log=all_symbols_strat.on_log)
  return all_symbols_strat.get_all_symbols()


def get_all_symbols_from_og_archive_from_flags():
  baseflags = get_base_flags()
  osflags = get_order_subscription_flags()
  return get_all_symbols_from_og_archive(baseflags, osflags)


def get_trading_dates(trading_date_str):
  if trading_date_str.find("-") >= 0:
    tdfrom, tdto = trading_date_str.split("-")
    return pandas.date_range(tdfrom, tdto).to_pydatetime()
  else:
    return [datetime.datetime.strptime(trading_date_str, '%Y%m%d')]


def get_flags_with_full_symbols(full_symbols_str):
  full_symbols = full_symbols_str.split(",")
  for full_symbol in full_symbols:
    exchange, market_type, symbol = full_symbol.split(":")
    yield {'exchange': exchange, 'market_type': market_type, 'symbols': symbol}


def set_fs_dict(fs_dict):
  for key, value in (fs_dict or {}).items():
    setattr(flags.FLAGS, key, value)


def iter_flags_with_full_symbols(full_symbols_str):
  for flag_dict in get_flags_with_full_symbols(full_symbols_str):
    set_fs_dict(flag_dict)
    yield


def relative_norm(product):
  return (product.subscription_symbol
          if hasattr(product, 'subscription_symbol') else product.symbol)


class SerializedJob(object):
  def __init__(self, filename):
    self.filename = filename

  def deserialize(self):
    if not os.path.exists(self.filename):
      self.serialize(self.filename)
    return self.deserialize_impl(self.filename)

  def serialize(self, filename):
    raise NotImplementedError()

  def deserialize_impl(self, filename):
    return
